package ai.people.core.nebula.holder;

import ai.people.core.nebula.entity.NebulaTransaction;
import com.vesoft.nebula.client.graph.data.ResultSet;
import com.vesoft.nebula.client.graph.exception.IOErrorException;
import com.vesoft.nebula.client.graph.net.Session;
import com.vesoft.nebula.client.graph.net.SessionWrapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.NamedThreadLocal;

/**
 * nebula会话上下文持有人
 *
 * @author yuanqinglong
 * @date 2022/4/28 15:02
 */
@Slf4j
public class NebulaSessionContextHolder {

    private static final ThreadLocal<NebulaTransaction> requestSessionThreadHolder = new NamedThreadLocal<>("Nebula NebulaTransaction attributes");

    public static void setTx(NebulaTransaction tx) {
        requestSessionThreadHolder.set(tx);
    }

    public static NebulaTransaction getTx() {
        return requestSessionThreadHolder.get();
    }


    public static void resetHolder() {
        requestSessionThreadHolder.remove();
    }

    /**
     * 切换space
     *
     * @param space          空间
     * @param sessionWrapper 会话包装
     * @throws IOErrorException ioerror异常那么
     */
    public static void switchSpace(String space, SessionWrapper sessionWrapper) throws IOErrorException {
        doSwitch(space, sessionWrapper);
    }

    public static void switchSpace(String space, Session session) throws IOErrorException {
        doSwitch(space, session);
    }

    private static void doSwitch(String space, SessionWrapper sessionWrapper) throws IOErrorException {
        ResultSet resultSet = sessionWrapper.execute("USE " + space);
        if (!resultSet.isSucceeded()) {
            throw new RuntimeException(
                    "Switch space `"
                            + space
                            + "' failed: "
                            + resultSet.getErrorMessage());
        }
        log.info("nebula graph switch space to {}", space);
    }

    private static void doSwitch(String space, Session session) throws IOErrorException {
        ResultSet resultSet = session.execute("USE " + space);
        if (!resultSet.isSucceeded()) {
            throw new RuntimeException(
                    "Switch space `"
                            + space
                            + "' failed: "
                            + resultSet.getErrorMessage());
        }
    }

}
